#!/usr/bin/env python
# -*- coding:utf-8 -*-
'''
@File   :   mgf.py
@Author :   Song
@Time   :   2022/3/3 14:27
@Contact:   songjian@westlake.edu.cn
@intro  : 
'''
import numpy as np
import pandas as pd
import click
from pathlib import Path
import multiprocessing as mp
import pyteomics.mgf
import operator

import predifine

class mgf_Reader():

    def __init__(self, fpath, is_annotated=True):
        self.ms = None
        self.all_rt = np.array([])  # 表示每一次采集的rt
        self.TimeUnit = None
        self.mgf_is_annotated = is_annotated

        # type
        fpath = str(fpath)
        self.suffix = fpath.split('.')[-1].lower()

        self.load(fpath)
        self.init()

    def size(self):
        return len(self.ms)

    def load(self, file_path):
        assert self.suffix == 'mgf', 'Only .mgf is supported!'
        self.ms = pyteomics.mgf.read(file_path, convert_arrays=1, read_charges=False)

    def init(self):
        self.get_time_unit()
        self.__load_to_memory()

    def get_mz_suffix(self):
        return self.suffix

    def get_time_unit(self):
        if self.suffix == 'mzxml':
            if self.ms[self.size() - 1]['retentionTime'] < 3 * 60:
                self.TimeUnit = 'minute'
            else:
                self.TimeUnit = 'second'
        elif self.suffix == 'mzml':
            if self.ms[self.size() - 1]['scanList']['scan'][0]['scan start time'] < 5 * 60:
                self.TimeUnit = 'minute'
            else:
                self.TimeUnit = 'second'
        elif self.suffix == 'mgf':
            try:
                self.ms[0]['params']['rtinseconds']
                self.TimeUnit = 'second'
            except:
                self.TimeUnit = 'minute'

    def process_worker(self, idx_start, idx_end): # 作为进程worker不能私有
        pr_mz_v, pr_charge_v, pr_seq_v = [], [], []
        scan_x_v, scan_y_v, scan_len_v = [], [], []
        scan_NO_v = [] # mgf 自带的谱图编号
        if self.mgf_is_annotated:
            for idx in range(idx_start, idx_end):
                scan = self.ms[idx]
                pr_mz_v.append(scan['params']['pepmass'][0])
                pr_charge_v.append(scan['params']['charge'][0])
                pr_seq_v.append(scan['params']['seq'])
                x = scan['m/z array']
                y = scan['intensity array']
                y = y/y.max()
                scan_x_v.append(x)
                scan_len_v.append(len(x))
                scan_y_v.append(y)
            return {'pr_mz': pr_mz_v,
                    'pr_charge': pr_charge_v,
                    'pr_seq': pr_seq_v,
                    'scan_x': scan_x_v,
                    'scan_y': scan_y_v,
                    'scan_len': scan_len_v}
        else:
            for idx in range(idx_start, idx_end):
                scan = self.ms[idx]
                No = int(scan['params']['title'].split('.')[-2])
                scan_NO_v.append(No)
                pr_mz_v.append(scan['params']['pepmass'][0])
                pr_charge_v.append(scan['params']['charge'][0])
                x = scan['m/z array']
                y = scan['intensity array']
                y = y / y.max()
                scan_x_v.append(x)
                scan_len_v.append(len(x))
                scan_y_v.append(y)
            return {'pr_mz': pr_mz_v,
                    'pr_charge': pr_charge_v,
                    'scan_x': scan_x_v,
                    'scan_y': scan_y_v,
                    'scan_len': scan_len_v,
                    'scan_No': scan_NO_v}

    def __load_to_memory(self):
        self.scan_x, self.scan_y, self.scan_len = [], [], []
        self.pr_mz, self.pr_charge, self.pr_seq = [], [], []
        self.scan_No = []

        # multiprocess load data
        cpu_num = mp.cpu_count()
        process_num = int(cpu_num / 2)  # default cores / 2
        process_num = 8 if process_num >= 8 else process_num
        pool = mp.Pool(process_num)
        slices = np.ceil(np.linspace(0, len(self.ms), process_num+1)).astype(int)

        results = [pool.apply_async(self.process_worker, args=(slices[i], slices[i+1])) for i in range(process_num)]
        results = [r.get() for r in results] # get同步进程。result是四个字典
        pool.close()
        pool.join()

        for result in results:
            self.pr_mz.extend(result['pr_mz'])
            self.pr_charge.extend(result['pr_charge'])
            if self.mgf_is_annotated:
                self.pr_seq.extend(result['pr_seq'])
            self.scan_x.extend(result['scan_x'])
            self.scan_y.extend(result['scan_y'])
            self.scan_len.extend(result['scan_len'])
            self.scan_No.extend(result['scan_No'])

        self.pr_mz = np.array(self.pr_mz)
        self.pr_charge = np.array(self.pr_charge)
        self.scan_x = np.concatenate(self.scan_x)
        self.scan_y = np.concatenate(self.scan_y)
        self.scan_len = np.array(self.scan_len)
        self.scan_No = np.array(self.scan_No)

        if self.mgf_is_annotated:
            s = pd.Series(self.pr_seq)
            s = s.replace(['C\+57.021', 'M\+15.995'], ['c', 'm'], regex=True)
            self.seq_len = s.str.len().to_numpy()
            s = list(s.str.cat())
            f = operator.itemgetter(*s)
            seq_idx = f(predifine.g_aa_to_idx)
            self.seq_idx = np.array(seq_idx)